#include <stdio.h>
#include <stdlib.h>
#include <getopt.h>
#include <limits.h>
#include <string.h>
#include <stdint.h>
#include <openssl/ec.h>
#include <openssl/evp.h>
#include <openssl/sm2.h>
#include <openssl/sm3.h>
#include <openssl/bn.h>

#include "utils.h"
#include "ovmf_hash_table.h"
#include "session_data.h"
#include "cal_vm_digest.h"

typedef enum option_choice {
    OPT_ERR = -1,
    OPT_EOF = 0,
    OPT_HELP,
    OPT_BIOS,
    OPT_KERNEL,
    OPT_CMDLINE,
    OPT_INITRD,
    OPT_VCPU_MODEL,
    OPT_VCPU_NUM,
    OPT_CSV3,
    OPT_FAMILY,
    OPT_MODEL,
    OPT_STEPPING,
} option_choice_t;

char *vcpu_model[] = { "Dhyana", "Dhyana-v1", "Dhyana-v2", "Dhyana-v3", "Dharma", "Dharma-v1", "host"};

static struct option long_args_options[] = {
    {"help",          no_argument,       0, OPT_HELP},
    {"ovmf",          required_argument, 0, OPT_BIOS},
    {"kernel",        required_argument, 0, OPT_KERNEL},
    {"cmdline",       required_argument, 0, OPT_CMDLINE},
    {"initrd",        required_argument, 0, OPT_INITRD},
    {"csv3",          no_argument,       0, OPT_CSV3},
    {"vcpumodel",     required_argument, 0, OPT_VCPU_MODEL},
    {"smp",           required_argument, 0, OPT_VCPU_NUM},
    {"family",        required_argument, 0, OPT_FAMILY},
    {"model",         required_argument, 0, OPT_MODEL},
    {"stepping",      required_argument, 0, OPT_STEPPING},
    {0, 0, 0, 0}
};

/**
 * @brief  Check if the vcpu model is in the range
 *
 * @param  [in]  model     The vcpu model string
 *
 * @returns 0 for success, -1 for failure
 */
int good_vcpu_model(char *model)
{
    int array_size = sizeof(vcpu_model) / sizeof(vcpu_model[0]);
    int n = 0;

    if (!model)
        return -1;

    for(n = 0; n < array_size; n++) {
        if (strcmp(model, vcpu_model[n]) == 0) {
            return 0;
        }
    }

    printf("Wrong cpu model, please choose vcpu model from :\n");
    for(n = 0; n < array_size; n++) {
        printf("%s ", vcpu_model[n]);
    }
    printf("\n");

    return -1;
}

/**
 * @brief  Usage
 *
 */

void usage(void)
{
    printf("Usage: program [options]\n");
    printf("Options:\n");
    printf("  --help,                         Display this summary\n");
    printf("  --ovmf <guest_bios_file>,       Input the virtual machine bios file\n");
    printf("  --kernel <guest_kernel_file>,   [optional] Input the virtual machine kernel file\n");
    printf("  --cmdline <guest_cmdline>,      [optional] Input the virtual machine command line file\n");
    printf("  --initrd <guest_initrd_file>,   [optional] Input the virtual machine initrd file\n");
    printf("  --cpu <vcpu_model_name>,        [optional] Hygon vCPU model name 'host' or (Dhyana{,-v1,-2,-v3},Dharma{,-v1}), required by CSV3, This argument is mutually exclusive with `--famliy $X --model $Y --stepping $Z\n");
    printf("  --smp <vcpu_num>,               [optional] The number of vCPUs of the guest, required by CSV3\n");
    printf("  --csv3,                         the virtual machine is csv3\n");
    printf("  --family <family>,              [optional] The family of vCPUs of the guest, required by CSV3\n");
    printf("  --model <vcpu_num>,             [optional] The model of vCPUs of the guest, required by CSV3\n");
    printf("  --stepping <vcpu_num>,          [optional] The stepping of vCPUs of the guest, required by CSV3\n");
}



int main(int argc, char *argv[])
{
    int opt;
    int option_index              = 0;
    char *endptr                  = NULL;
    char *bios                    = NULL;
    char *kernel                  = NULL;
    char *cmdline                 = NULL;
    char *initrd                  = NULL;
    char *vcpu_model_arg          = NULL;
    unsigned long vcpu_num        = 0;
    uint8_t vm_digest[HASH_SIZE]  = {0};
    guest_policy    policy        = {0};
    bool  is_csv3                 = 0;
    char *family_str              = NULL;
    char *model_str               = NULL;
    char *stepping_str            = NULL;
    uint64_t vcpu_model_value     = 0;
    char * base64_digest          = NULL;

    int ret                       = -1;

    if (argc <= 1) {
        usage();
        return -1;
    }

    while ((opt = getopt_long(argc, argv, "", long_args_options, &option_index)) != -1) {
        switch (opt) {
            case OPT_HELP:
                usage();
                return 0;
            case OPT_BIOS:
                if (optarg == NULL || optarg[0] == '-') {
                    usage();
                    return -1;
                }
                bios = optarg;
                printf("bios file = %s\n", bios);
                break;
            case OPT_KERNEL:
                kernel = optarg;
                printf("kernel file = %s\n", kernel);
                break;
            case OPT_CMDLINE:
                cmdline = optarg;
                printf("cmdline = '%s'\n", cmdline);
                break;
            case OPT_INITRD:
                initrd = optarg;
                printf("initrd = %s \n", initrd);
                break;
            case OPT_CSV3:
                is_csv3 = 1;
                break;
            case OPT_VCPU_MODEL:
                vcpu_model_arg = optarg;
                if (good_vcpu_model(vcpu_model_arg) == -1)
                    return -1;
                printf("vcpu_model = %s \n", vcpu_model_arg);
               break;
            case OPT_VCPU_NUM:
                vcpu_num = strtoul(optarg, &endptr, 10);
                if (*endptr != '\0') {
                    printf("Wrong vcpu num: %s\n", optarg);
                    return -1;
                }
                printf("vcpu num = %lu \n", vcpu_num);
                break;
            case OPT_FAMILY:
                family_str = optarg;
                printf("family = %s \n", family_str);
                break;
            case OPT_MODEL:
                model_str = optarg;
                printf("model = %s \n", model_str);
                break;
            case OPT_STEPPING:
                stepping_str = optarg;
                printf("stepping = %s \n", stepping_str);
                break;
            default:
                usage();
                return -1;
        }
    }


    if (is_csv3)
        policy.csv3 = 1;

    if (is_csv3) {
        if (vcpu_model_arg != NULL) {
            if (family_str != NULL || model_str != NULL || stepping_str != NULL) {
                printf("vcpu_mode is mutually exclusive with `--famliy $X --model $Y --stepping $Z\n");
                return -1;
            }

            vcpu_model_value = get_vcpu_model_value(vcpu_model_arg);
            if (!vcpu_model_value){
                printf("vcpu model is 0, invalid\n");
                return -1;
            }
        } else {
            uint32_t family, model, stepping;

            if (family_str == NULL || model_str == NULL || stepping_str == NULL) {
                printf("invalid family/model/stepping\n");
                return -1;
            }
            family = str2int(family_str);
            model  = str2int(model_str);
            stepping = str2int(stepping_str);

            vcpu_model_value = build_vcpu_model(family, model, stepping);
        }
    }

    if (calculate_vm_digest( policy, bios, kernel, cmdline, initrd,
                            vcpu_model_value, vcpu_num,
                            vm_digest) != 0) {
        printf("fail to calculate VM digest\n");
        goto finish;
    }

    ret = 0;
    printf("Hex Digest:\n");
    hex_dump(vm_digest, sizeof(vm_digest));

    base64_digest = base64_encode(vm_digest, sizeof(vm_digest));
    if (base64_digest) {
        printf("Base64 Digest:\n");
        printf("%s\n", base64_digest);
        free(base64_digest);
    }

finish:

    return ret;
}
